"""
model name : MY_YOLO
file       : show.py
information:
    author : OuYang
    time   : 2025/1/22
"""
import cv2
import numpy as np


def PIL_to_cv2(image):
    image = np.array(image)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image


def tensorImage_to_cv2(image):
    image = np.array(image)
    image = np.transpose(image, (1, 2, 0))
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    return image


def show_image(image):
    cv2.imshow('image', image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def show_image_from_PIL(image):
    image = PIL_to_cv2(image)
    show_image(image)


def show_image_from_tensor(image):
    image = tensorImage_to_cv2(image)
    show_image(image)


def draw_boxes(image, boxes, texts=None, scares=None, color=(0, 255, 0), thickness=1, show=False):
    """
    在图片上绘制矩形框
        Note: 在原图上绘制
    :param image: 待绘制的图片，在原图上绘制
    :param boxes: 需要绘制的所有框
    :returns 绘制好的图像
    """
    for i, box in enumerate(boxes):
        image = cv2.rectangle(
            image,
            (int(box[0]), int(box[1])),
            (int(box[2]), int(box[3])),
            color,
            thickness
        )
        if texts is not None and scares is not None:
            str_text = str(texts[i] + " " + str(round(scares[i], 2)))
            (text_w, text_h), baseline = cv2.getTextSize(str_text, cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, thickness=thickness)
            top_left = (int(box[0]), int(box[1]) - text_h - baseline)
            bottom_right = (int(box[0]) + text_w, int(box[1]))
            cv2.rectangle(image, top_left, bottom_right, color, thickness=-1)
            cv2.putText(
                image,
                str_text,
                (int(box[0]), int(box[1]) - 5),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (255, 255, 255),
                thickness=1
            )

    if show:
        show_image(image)
    return image


def draw_boxes_from_PIL(image, boxes, texts=None, scares=None, color=(0, 255, 0), thickness=1, show=False):
    image = PIL_to_cv2(image)
    draw_boxes(image, boxes, texts, scares, color, thickness, show)
    return image


def draw_boxes_from_tensor(image, boxes, texts=None, scares=None, color=(0, 255, 0), thickness=1, show=False):
    image = tensorImage_to_cv2(image)
    draw_boxes(image, boxes, texts, scares, color, thickness, show)
    return image
